import rsa
import time

print('''
    easyRSA - It just encrypt and decrypt
    Copyright (C) 2021  stak
    
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
''')

def gen_rsa(bit=1024):
    global pubkey, privkey
    if __name__ == '__main__':
        bit = input('Enter number of bits for generate RSA keys without spaces(default 1024): ')
        if bit == '':
            bit = 1024
        bit = int(bit)
    (pubkey, privkey) = rsa.newkeys(bit)
    pubkey_pem = pubkey.save_pkcs1()
    private_pem = privkey.save_pkcs1()

def encryption(word):
    global encrypted_message
    word = word.encode('utf-8')
    encrypted_message = rsa.encrypt(word, pubkey)
    return encrypted_message

def unencryption(encrypted_message):
    unencrypted_message = rsa.decrypt(encrypted_message, privkey)
    global decode_message
    decode_message = unencrypted_message.decode('utf-8')
    return decode_message

def console_picture():
    print('                      ____  ____    _     ')
    time.sleep(0.1)
    print('  ___  __ _ ___ _   _|  _ \/ ___|  / \    ')
    time.sleep(0.1)
    print(' / _ \/ _` / __| | | | |_) \___ \ / _ \   ')
    time.sleep(0.1)
    print('|  __/ (_| \__ \ |_| |  _ < ___) / ___ \  ')
    time.sleep(0.1)
    print(' \___|\__,_|___/\__, |_| \_\____/_/   \_\ ')
    time.sleep(0.1)
    print('                |___/                     ')

def control():
    exit_words = ['exit', 'stop', 'break', '/exit', '/stop', '/break']
    while True:
        instruction = input('Enter instruction: ')
        if instruction == '/start':
            gen_rsa()
            message = input('Enter word: ')
            encryption(message)
            print(f'encrypted message = {encrypted_message}')
        if instruction == '/e':
            try:
                word = input('Enter word: ')
                encryption(word)
                print(f'encrypted message = {encrypted_message}')
            except NameError:
                print("You didn't generate RSA keys")
                userinput = input('Generate RSA keys?[y,n]: ')
                if userinput == 'y':
                    gen_rsa()
                else:
                    print('RSA keys have not been generated!')
                    continue

            finally:
                continue

        if instruction == '/d':
            try:
                unencryption(encrypted_message=encrypted_message)
                print(f'decode_message = {decode_message}')
            except NameError:
                print("You didn't encrypt any words")
            finally:
                continue

        if instruction == '/g':
            gen_rsa()

        if instruction in exit_words:
            console_picture()
            print('Bye!')
            break

def main():
    console_picture()
    control()

if __name__ == "__main__":
    main()
